from PIL import Image
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms

# 1、transform如何使用
# 2、为什么需要tensor数据类型
# 包装了很多神经网络需要的字段，比如反向传播函数，步长等
img = Image.open('../data/BEES_ANTS/train/ants/0013035.jpg')
trans_to_tensor = transforms.ToTensor()
img_tensor = trans_to_tensor(img)

print(img_tensor.shape)

# ToTensor
writer = SummaryWriter("../logs")
writer.add_image('image', img_tensor, 1)

# Normalize
trans_norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
img_tensor = trans_norm(img_tensor)
writer.add_image('image', img_tensor, 2)

trans_norm = transforms.Normalize(mean=[0.6, 0.5, 0.4], std=[0.3, 0.2, 0.1])
img_tensor = trans_norm(img_tensor)
writer.add_image('image', img_tensor, 3)

# resize 如果只给一个Int值，就按比例把窄边变换成指定的数
trans_resize = transforms.Resize((512, 512))
img_tensor = trans_resize(img_tensor)
# print(img_tensor)
# img_tensor = trans_to_tensor(img_tensor)
writer.add_image('image', img_tensor, 4)

# Compose
trans_compose = transforms.Compose([transforms.Resize(1014), transforms.ToTensor()])
img_tensor = trans_compose(img)
writer.add_image('image', img_tensor, 5)

# randomCrop
trans_random_crop = transforms.RandomCrop((300, 300))
trans_compose2 = transforms.Compose([trans_random_crop, transforms.ToTensor()])
for i in range(10):
    img_tensor = trans_compose2(img)
    writer.add_image('image', img_tensor, i+6)
writer.close()
